# app.py
import cv2
import mediapipe as mp
import numpy as np
from flask import Flask, Response, render_template_string, request, jsonify
import threading
import os

app = Flask(__name__)

# --------- 全局变量（线程安全访问） ----------
lock = threading.Lock()
current_glasses_name = "glasses1.png"  # 默认
GLASSES_FOLDER = os.path.join("static", "glasses")

# 加载所有可用眼镜文件名（放在 static/glasses/ 目录下）
available_glasses = []
if os.path.isdir(GLASSES_FOLDER):
    for f in os.listdir(GLASSES_FOLDER):
        if f.lower().endswith((".png", ".webp")):
            available_glasses.append(f)
else:
    os.makedirs(GLASSES_FOLDER, exist_ok=True)

if not available_glasses:
    # 提醒用户放文件
    print("提示：把透明背景的眼镜 PNG 放到 static/glasses/ 目录下，命名为 glasses1.png 等。")

# --------- MediaPipe 初始化 ----------
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(
    max_num_faces=1, refine_landmarks=True,
    min_detection_confidence=0.5, min_tracking_confidence=0.5
)

# 打开摄像头（0 默认），若要在服务器上运行并使用文件或流，请修改这里
cap = cv2.VideoCapture(0)

# --------- 帮助函数：根据名字加载眼镜（返回 None 则表示无） ----------
def load_glasses_by_name(name):
    path = os.path.join(GLASSES_FOLDER, name)
    if os.path.exists(path):
        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        if img is None:
            return None
        return img
    return None

# --------- MJPEG 生成器 ----------
def gen_frames():
    global current_glasses_name
    while True:
        success, frame = cap.read()
        if not success:
            break

        # 处理帧：MediaPipe 检测 + 合成当前眼镜
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = face_mesh.process(rgb_frame)

        # 读取当前眼镜（每帧重新加载以便实时切换；可改为缓存）
        with lock:
            gname = current_glasses_name
        overlay_img = None
        if gname:
            overlay_img = load_glasses_by_name(gname)

        if results.multi_face_landmarks and overlay_img is not None:
            for face_landmarks in results.multi_face_landmarks:
                h, w, _ = frame.shape
                # 眼角点：33 左眼内角，263 右眼内角（参考 FaceMesh 索引）
                left = face_landmarks.landmark[33]
                right = face_landmarks.landmark[263]
                lx, ly = int(left.x * w), int(left.y * h)
                rx, ry = int(right.x * w), int(right.y * h)

                # 计算眼镜宽度（两眼间距 *系数）
                eye_dist = int(((rx - lx) ** 2 + (ry - ly) ** 2) ** 0.5)
                eye_width = max(1, int(eye_dist * 2.0))  # 防止为0

                # 缩放眼镜
                gh, gw = overlay_img.shape[0], overlay_img.shape[1]
                scale = eye_width / float(gw)
                new_w = eye_width
                new_h = max(1, int(gh * scale))
                resized = cv2.resize(overlay_img, (new_w, new_h), interpolation=cv2.INTER_AREA)

                # 放在双眼中点略上方
                cx, cy = (lx + rx) // 2, (ly + ry) // 2
                x1 = cx - new_w // 2
                y1 = cy - int(new_h * 0.45)  # 微调使眼镜略遮住眼睛
                x2, y2 = x1 + new_w, y1 + new_h

                # 边界裁剪
                H, W, _ = frame.shape
                if x1 < 0:
                    x_offset = -x1
                    x1 = 0
                else:
                    x_offset = 0
                if y1 < 0:
                    y_offset = -y1
                    y1 = 0
                else:
                    y_offset = 0

                x2 = min(W, x2)
                y2 = min(H, y2)

                # 计算对应 resized 区域
                rw = x2 - x1
                rh = y2 - y1
                if rw <= 0 or rh <= 0:
                    continue

                resized_crop = resized[y_offset:y_offset+rh, x_offset:x_offset+rw]

                # 分离 alpha
                if resized_crop.shape[2] == 4:
                    overlay_rgb = resized_crop[:, :, :3]
                    overlay_alpha = resized_crop[:, :, 3] / 255.0
                    overlay_alpha = overlay_alpha[..., None]
                    roi = frame[y1:y1+rh, x1:x1+rw]
                    # 合成
                    try:
                        frame[y1:y1+rh, x1:x1+rw] = (overlay_alpha * overlay_rgb + (1 - overlay_alpha) * roi).astype(np.uint8)
                    except Exception as e:
                        # 防止索引问题
                        pass
                else:
                    # 没有 alpha 通道，直接覆盖（不推荐）
                    frame[y1:y1+rh, x1:x1+rw] = resized_crop[:rh, :rw]

        # 编码为 JPEG
        ret, buffer = cv2.imencode('.jpg', frame)
        if not ret:
            continue
        frame_bytes = buffer.tobytes()

        # 生成 multipart frame
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

# --------- Flask 路由 ----------
INDEX_HTML = """
<!doctype html>
<html>
<head>
  <meta charset="utf-8">
  <title>虚拟试戴（Web Demo）</title>
  <style>
    body{font-family: Arial; text-align:center; background:#222; color:#fff}
    #video{border:4px solid #333; margin-top:10px}
    .controls{margin-top:10px}
    select, button{padding:8px 12px; font-size:16px}
  </style>
</head>
<body>
  <h1>虚拟试戴 — Web Demo</h1>
  <div class="controls">
    <label>选择眼镜：</label>
    <select id="glassesSelect"></select>
    <button onclick="apply()">应用</button>
  </div>
  <div>
    <img id="video" src="/video_feed" width="720" />
  </div>
<script>
async function loadList(){
  const resp = await fetch('/glasses_list');
  const data = await resp.json();
  const sel = document.getElementById('glassesSelect');
  sel.innerHTML = '';
  data.forEach(name=>{
    const opt = document.createElement('option');
    opt.value = name;
    opt.text = name;
    sel.appendChild(opt);
  });
}
async function apply(){
  const sel = document.getElementById('glassesSelect');
  const name = sel.value;
  await fetch('/set_glasses', {method:'POST', headers:{'Content-Type':'application/json'}, body: JSON.stringify({name})});
  // 轻刷新 img，防缓存
  const img = document.getElementById('video');
  img.src = '/video_feed?rand=' + Math.random();
}
window.onload = loadList;
</script>
</body>
</html>
"""

@app.route('/')
def index():
    return render_template_string(INDEX_HTML)

@app.route('/video_feed')
def video_feed():
    return Response(gen_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/glasses_list')
def glasses_list():
    return jsonify(available_glasses)

@app.route('/set_glasses', methods=['POST'])
def set_glasses():
    global current_glasses_name
    data = request.get_json()
    name = data.get('name') if data else None
    if not name:
        return jsonify({"ok": False, "msg": "no name"}), 400
    if name not in available_glasses:
        return jsonify({"ok": False, "msg": "not found"}), 404
    with lock:
        current_glasses_name = name
    return jsonify({"ok": True, "current": current_glasses_name})

# 优雅关闭摄像头（若需要）
@app.route('/shutdown_cam')
def shutdown_cam():
    try:
        cap.release()
    except:
        pass
    return "camera released"

if __name__ == "__main__":
    # 生产环境请用 gunicorn/uvicorn 等 WSGI 服务器并注意摄像头权限
    app.run(host="0.0.0.0", port=5000, debug=True)